from langchain_ollama import ChatOllama
from langchain_ollama import OllamaEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from flask_cors import CORS  # 导入 CORS
from flask import Flask, request, Response
import json

from langchain_community.document_loaders.csv_loader import CSVLoader

app = Flask(__name__)
CORS(app)  # 启用 CORS

# 初始化 Ollama 模型和嵌入
model = "qwen2.5:14b-instruct-q3_K_S"
llm = ChatOllama(model=model, temperature=0.6)
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# 准备文档
loader = CSVLoader(
    'E:/罗湖项目/教育信息.csv'
)
docs = loader.load()  # 返回的是 Document 对象列表
texts = [doc.page_content for doc in docs]  # 提取所有文档的文本内容
text = ' '.join(texts)  # 使用空格将数组中的字符串元素连接成一个字符串

print(text)
 
# 分割文本
text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
chunks = text_splitter.split_documents(docs)  # 使用 split_documents 方法
chunk_texts = [chunk.page_content for chunk in chunks]  # 提取所有分块的文本内容

# 创建向量存储
vectorstore = FAISS.from_texts(chunk_texts, embeddings)  # 传入文本列表
retriever = vectorstore.as_retriever()

# 创建提示模板
template = """根据以下内容进行思考、总结，并基于文档内容来回答问题:
{context}

问题: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

# 创建检索-问答链
chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
)

# 定义流式响应函数
def generate_response(question):
    try:
        responseStr = '';
        for chunk in chain.stream(question):
            try:
                # 尝试解析为 JSON
                if isinstance(chunk, dict):
                    content = chunk.get('content', chunk)
                else:
                    content = str(chunk)
                
                # 确保换行符被保留
                content = content.replace('\\n', '\n')
                content = content.replace('\\r', '\n')

                arr = content.split(" ")
                content = arr[0].split("=")[1]
                content = content.replace('\'', "")
                
                responseStr += content
                # print(f"{content}")
                
                if(content) :
                    # 使用 JSON 编码确保特殊字符被正确处理
                    formatted_content = json.dumps({"model":f"{model}","created_at":"","message":{"role":"assistant","content":f"{content}"},"done":False})
                    yield f"{formatted_content}\n\n"                 

            except Exception as e:
                print(f"1-异常：{str(e)}")
                # 如果 JSON 解析失败，直接返回原始内容
                yield f"data: {json.dumps({'content': str(chunk)})}\n\n"
        print(responseStr)
            
    except Exception as e:
        print(f"2-异常：{str(e)}")
        yield f"data: {json.dumps({'error': str(e)})}\n\n"

@app.route('/ask', methods=['POST'])
def ask():
    data = request.get_json()
    question = data.get('messages', '')
    
    if not question:
        return {"status": "error", "message": "问题不能为空"}, 400
    
    return Response(generate_response(question), mimetype='text/event-stream')

if __name__ == "__main__":
    app.run(host='127.0.0.1', port=8002) 